from src.LoadData.dataLoader import dataclean, Voc, getWordIndex
from src.Main.config import MAX_LENGTH
import os
import torch
import random

QINGYUN_PATH = "../../Datasets/RawData/qingyun_seg"
save_dir = "../../Datasets/FinalData/"


def getQYPairs(qingyun_path):
    """
    :param qingyun_path: 青云数据集路径
    :return:
    """
    pairs = []
    with open(qingyun_path, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            pair = [[], []]
            pair[0] = dataclean(line.split(" | ")[0].strip()).strip()
            pair[1] = dataclean(line.split(" | ")[1].strip()).strip()
            length1 = len(pair[0].split(" "))
            length2 = len(pair[1].split(" "))
            if length1 > MAX_LENGTH or length2 > MAX_LENGTH or length1 < 2 or length2 < 2:
                pass
            else:
                pairs.append(pair)
    return pairs


def getTrainTestPairs(pairs, scale=0.8):
    """
    划分测试集和训练集
    :param pairs:原始数据集
    :param scale:比例
    :return:train, test
    """
    train_data_set_len = int(len(pairs) * scale)
    random.shuffle(pairs)
    return pairs[:train_data_set_len], pairs[train_data_set_len:]


if __name__ == '__main__':
    print("-----loading pairs----")
    pairs = getQYPairs(qingyun_path=QINGYUN_PATH)
    print(len(pairs))
    print("-----loading voc------")
    voc = Voc("qingyun")
    for pair in pairs:
        voc.addSentence(pair[0])
        voc.addSentence(pair[1])
    voc.trim(min_count=3)
    print(voc.num_words)
    print('' in voc.word2index)

    print("-----split train/test------")
    train_pairs, test_pairs = getTrainTestPairs(pairs, scale=0.8)

    print("-----loading index pairs------")
    train_pairs = getWordIndex(voc, train_pairs)
    test_pairs = getWordIndex(voc, test_pairs)

    print("------------save--------------")
    directory = os.path.join(save_dir, 'qingyun_data', "qingyun")
    if not os.path.exists(directory):
        os.makedirs(directory)
    torch.save(voc, os.path.join(directory, '{!s}.tar'.format('voc_qingyun')))
    torch.save(train_pairs, os.path.join(directory, '{!s}.tar'.format('pairs_qingyun_train')))
    torch.save(test_pairs, os.path.join(directory, '{!s}.tar'.format('pairs_qingyun_test')))
